import json
from agents import req_analyzer, usecase_generator, class_generator
from agents import sequence_generator, ocl_generator, model_evaluator
from dsl_definitions import *

def parse_json_safely(json_str: str) -> dict:
    """增强型JSON解析器，带详细错误诊断"""
    content = json_str.strip()
    
    # 处理代码块标记
    if content.startswith("```json") and content.endswith("```"):
        content = content[len("```json"):-len("```")].strip()
    elif content.startswith("```") and content.endswith("```"):
        content = content[len("```"):-len("```")].strip()
    
    # 尝试解析
    try:
        return json.loads(content)
    except json.JSONDecodeError as e:
        # 详细错误报告
        error_info = f"""
        ❌ JSON解析失败：
        - 错误位置：{e.pos}
        - 错误行：{e.lineno}, 列：{e.colno}
        - 问题片段：{content[e.pos-20:e.pos+20] if e.pos>20 else content[:40]}
        - 完整原始内容（前500字符）：{content[:500]}
        """
        return {"error": error_info}

def print_stage_header(stage: str):
    """打印阶段标题"""
    print(f"\n{'='*20} {stage} {'='*20}")

def optimize_class(class_json, feedback):
    """根据反馈优化类图"""
    prompt = f"""
    根据以下反馈优化类图：
    {feedback}
    当前类图：{json.dumps(class_json)}
    
    要求：
    1. 补全缺失的类和方法
    2. 调整关联关系
    3. 保持JSON格式
    """
    optimized_dsl = class_generator(prompt)
    return parse_json_safely(optimized_dsl)

def optimize_sequence(sequence_json, class_json, usecase_json):
    """根据类图和用例图优化顺序图"""
    prompt = f"""
    根据以下内容优化顺序图：
    现有顺序图：{json.dumps(sequence_json)}
    最新类图：{json.dumps(class_json)}
    用例图：{json.dumps(usecase_json)}
    
    要求：
    1. 确保消息匹配最新类方法
    2. 保持用例流程逻辑
    """
    optimized_dsl = sequence_generator(json.dumps(usecase_json), json.dumps(class_json))
    return parse_json_safely(optimized_dsl)

def optimize_ocl(ocl_json, class_json):
    """根据类图优化OCL合约"""
    prompt = f"""
    根据以下内容优化OCL合约：
    现有OCL：{json.dumps(ocl_json)}
    最新类图：{json.dumps(class_json)}
    
    要求：
    1. 确保OCL条件匹配最新类结构
    2. 为新增方法添加OCL约束
    """
    optimized_dsl = ocl_generator(json.dumps(class_json))
    return parse_json_safely(optimized_dsl)

def main():
    # 用户需求
    user_prompt = (
        "设计一个覆盖用户端与管理后台的共享单车系统，包含：\n"
        "- 用户注册/登录、实名认证、押金管理（支持信用免押）；\n"
        "- 扫码解锁车辆，处理异常场景（车辆故障、二维码损坏、定位失败）；\n"
        "- 行程实时数据（时长、里程、卡路里）、暂停用车（临时锁车）；\n"
        "- 结束用车自动计费，支持微信/支付宝/余额支付，及发票申请；\n"
        "- 会员月卡/季卡（免押金、骑行折扣、免费时长）；\n"
        "- 管理后台：车辆投放调度（热力图）、故障维修工单、数据统计；\n"
        "- 非功能需求：高并发（万级同时扫码）、定位精度≤10米、数据加密存储。"
    )

    # ========== 阶段1：需求分析 ==========
    print_stage_header("阶段1：需求分析")
    req_result = req_analyzer(user_prompt)
    print(f"📤 需求分析Agent原始响应：\n{req_result}\n")  # 强制打印原始响应
    req_json = parse_json_safely(req_result)
    
    if "error" in req_json:
        print(f"❌ 需求分析失败：{req_json['error']}")
        return
    
    # 打印需求分析结果
    print(f"✅ 需求分析成功：{len(req_json.get('functional', []))}个功能需求，{len(req_json.get('non_functional', []))}个非功能需求")

    # ========== 阶段2：生成用例图 ==========
    print_stage_header("阶段2：生成用例图")
    usecase_dsl = usecase_generator(json.dumps(req_json))
    print(f"📤 用例图Agent原始响应：\n{usecase_dsl}\n")  # 强制打印原始响应
    usecase_json = parse_json_safely(usecase_dsl)
    
    if "error" in usecase_json:
        print(f"❌ 用例图生成失败：{usecase_json['error']}")
        return
    
    # 打印用例图结果
    print(f"✅ 用例图生成成功：{len(usecase_json.get('usecases', []))}个用例，{len(usecase_json.get('actors', []))}个参与者")

    # ========== 阶段3：生成类图 ==========
    print_stage_header("阶段3：生成类图")
    initial_class_prompt = f"根据用例图生成类图：{json.dumps(usecase_json)}"
    class_dsl = class_generator(initial_class_prompt)
    print(f"📤 类图Agent原始响应：\n{class_dsl}\n")  # 强制打印原始响应
    class_json = parse_json_safely(class_dsl)
    
    if "error" in class_json:
        print(f"❌ 类图生成失败：{class_json['error']}")
        return
    
    # 打印类图结果
    print(f"✅ 类图生成成功：{len(class_json.get('classes', []))}个类，{len(class_json.get('associations', []))}个关联")

    # ========== 阶段4：生成系统顺序图 ==========
    print_stage_header("阶段4：生成系统顺序图")
    sequence_dsl = sequence_generator(json.dumps(usecase_json), json.dumps(class_json))
    print(f"📤 顺序图Agent原始响应：\n{sequence_dsl}\n")  # 强制打印原始响应
    sequence_json = parse_json_safely(sequence_dsl)
    
    if "error" in sequence_json:
        print(f"❌ 顺序图生成失败：{sequence_json['error']}")
        return
    
    # 打印顺序图结果
    print(f"✅ 顺序图生成成功：{len(sequence_json.get('messages', []))}条消息，{len(sequence_json.get('participants', []))}个参与者")

    # ========== 阶段5：生成OCL合约 ==========
    print_stage_header("阶段5：生成OCL合约")
    ocl_dsl = ocl_generator(json.dumps(class_json))
    print(f"📤 OCL合约Agent原始响应：\n{ocl_dsl}\n")  # 强制打印原始响应
    ocl_json = parse_json_safely(ocl_dsl)
    
    if "error" in ocl_json:
        print(f"❌ OCL合约生成失败：{ocl_json['error']}")
        return
    
    # 打印OCL结果
    method_count = sum(len(cls.get('methods', [])) for cls in ocl_json.get('classes', []))
    print(f"✅ OCL合约生成成功：为{len(ocl_json.get('classes', []))}个类定义了{method_count}个方法约束")

    

if __name__ == "__main__":
    main()
